from transformers import AutoModelForCausalLM, AutoTokenizer
import time
from openai import OpenAI 
import os


# prompt-2 ，添加性别、年龄评估，其余同prompt1一致
system = r'''
你是一个小说的角色分类模型，可以根据输入的小说内容，完成以下任务： 将每位角色的台词单独提取，并在台词前添加角色名（若未明确角色名，则标记为“未知角色”），并评估角色的性别、年龄。 将所有非角色台词的部分标记为“解说”，并单独提取。 保持原文的语义完整，不遗漏任何信息。输出顺本遵循文本输入顺序。输出格式举例：[解说]：“”， [角色1_性别_年龄：]“”，[解说：]“”， [角色2_性别_年龄：]“”，等等。输入为：
'''
input = r'''
第五章 聚气散。 “咳。”白袍老者轻咳了一声，站起身来对着萧战拱了拱手，微笑道：“萧族长，此次前来贵家族，主要是有事相求！” “呵呵，葛叶先生，有事请说便是，如果力所能及，萧家应该不会推辞。”对于这位老者，萧战可不敢怠慢，连忙站起来客气的道，不过由于不知道对方到底所求何事，所以也不敢把话说得太满。 “呵呵，萧族长，你可认识她么？”葛叶微微一笑，指着身旁的少女含笑问道。 “呃…恕萧战眼拙，这位小姐…”闻言，萧战一愣，上下打量了一下少女，略微有些尴尬的摇了摇头。 当年纳兰嫣然被云韵收为弟子之时，年仅十岁，在云岚宗中修炼了五年时间，所谓女大十八变，好多年未见，萧战自然不知道面前的少女，便是自己名义上的儿媳妇。 “咳…她的名字叫纳兰嫣然。” “纳兰嫣然？纳兰老爷子的孙女纳兰嫣然？”萧战先是一怔，紧接着满脸大喜，想必是记起了当年的那事，当下，急忙对着少女露出温和的笑容：“原来是纳兰侄女，萧叔叔可有好多年未曾与你见面了，可别怪罪叔叔眼拙。” 忽然出现的一幕，让得众人也是略微一愣，三位长老互相对视了一眼，眉头不由得皱了皱… “萧叔叔，侄女一直未曾前来拜见，该赔罪的，可是我呢，哪敢怪罪萧叔叔。”纳兰嫣然甜甜的笑道。 “呵呵，纳兰侄女，以前便听说了你被云韵大人收入门下，当时还以为是流言，没想到，竟然是真的，侄女真是好天赋啊…”萧战笑着赞叹道。 “嫣然只是好运罢了…”浅浅一笑，纳兰嫣然有些吃不消萧战的热情，桌下的手掌，轻轻扯了扯身旁的葛叶。 “呵呵，萧族长，在下今日所请求之事，便与嫣然有关，而且此事，还是宗主大人亲自开口…”葛叶轻笑了一声，在提到宗主二字时，脸庞上的表情，略微郑重。 脸色微微一变，萧战也是收敛了笑容，云岚宗宗主云韵可是加玛帝国的大人物，他这小小的一族之长，可是半点都遭惹不起，可以她的实力与势力，又有何事需要萧家帮忙？葛叶说是与纳兰侄女有关，难道？ 想到某种可能，萧战的嘴角忍不住的抽搐了几下，硕大的手掌微微颤抖，不过好在有着袖子的遮掩，所以也未曾被发现，强行压下心头的怒火，声音有些发颤的凝声道：“葛叶先生，请说！” “咳…”葛叶脸色忽然出现了一抹尴尬，不过想起宗主对纳兰嫣然的疼爱，又只得咬了咬牙，笑道：“萧族长，您也知道，云岚宗门风严厉，而且宗主大人对嫣然的期望也是很高，现在基本上已经是把她当做云岚宗下一任的宗主在培养…而因为一些特殊的规矩，宗主传人在未成为正式宗主之前，都不可与男子有纠葛…” “宗主大人在询问过嫣然之后，知道她与萧家还有一门亲事，所以…所以宗主大人想请萧族长，能够…解除了这婚约。” “咔！”萧战手中的玉石杯，轰然间化为了一蓬粉末。 大厅之中，气氛有些寂静，上方的三位长老也是被葛叶的话震了了震，不过片刻之后，他们望向萧战的目光中，已经多出了一抹讥讽与嘲笑。 “嘿嘿，被人上门强行解除婚约，看你这族长，以后还有什么威望管理家族！” 一些年轻一辈的少年少女并不知晓萧炎与纳兰嫣然的婚约，不过在向身旁的父母打听了一下之后，他们的脸色，顿时变得精彩了起来，讥诮的嘲讽目光，投向了角落处的萧炎… 望着萧战那阴沉至极的脸色，纳兰嫣然也是不敢抬头，将头埋下，手指紧张的绞在了一起。 “萧族长，我知道这要求有些强人所难，不过还请看在宗主大人的面上，解除了婚约吧…”无奈的叹了一口气，葛叶淡淡的道。 萧战拳头紧握，淡淡的青色斗气，逐渐的覆盖了身躯，最后竟然隐隐约约的在脸庞处汇聚成了一个虚幻的狮头。 萧家顶级功法：狂狮怒罡！等级：玄阶中级！ 望着萧战的反映，葛叶脸庞也顿时凝重了起来，身体挡在纳兰嫣然身前，鹰爪般的双手猛的曲拢，青色斗气在鹰爪中汇聚而起，散发着细小而凌厉的剑气。 云岚宗高深功法，青木剑诀！等级：玄阶低级！ 随着两人气息的喷发，大厅之中，实力较弱的少年们，脸色猛的一白，旋即胸口有些发闷。 就在萧战的呼吸越加急促之声，三位长老的厉喝声，却是宛如惊雷般的在大厅中响起：“萧战，还不住手！你可不要忘记，你是萧家的族长！” 身子猛的一僵，萧战身体上的斗气缓缓的收敛，最后完全消失。 一屁股坐回椅子上，萧战脸色淡漠的望着低头不言的纳兰嫣然，声音有些嘶哑的道：“纳兰侄女呐，好魄力啊，纳兰肃有你这女儿，真是很让人羡慕啊！” 娇躯微微一颤，纳兰嫣然呐呐的道：“萧叔叔…” “呵呵，叫我萧族长就好，叔叔这称谓，我担不起，你是未来云岚宗的宗主，日后也是斗气大陆的风云人物，我家炎儿不过是资质平庸之辈，也的确是配不上你…”淡淡的挥了挥手，萧战语气冷漠的道。 “多谢萧族长体谅了。”闻言，一旁的葛叶大喜，对着萧战赔笑道:“萧族长，宗主大人知道今天这要求很是有些不礼貌，所以特地让在下带来一物，就当做是赔礼！” 说着，葛叶伸手抹了抹手指上的一枚戒指，一只通体泛绿的古玉盒子在手中凭空出现… 小心的打开盒子，一股异香顿时弥漫了大厅，闻者皆都是精神为之一畅。 三位长老好奇的伸过头，望着玉匣子内，身体猛的一震，惊声道：“聚气散？”
'''

# -- 输出划分后保存目录 --
path_save_txt = './out_story_txt/'
# 自定义保存章节名
save_file = path_save_txt+'3.txt'
os.makedirs(path_save_txt, exist_ok=True)

save_file = open(save_file, 'w', encoding='utf-8')

client = OpenAI(
    api_key="YOUR API KEY",
    base_url="https://open.bigmodel.cn/api/paas/v4/"
) 

completion = client.chat.completions.create(
    model="glm-4-plus",  # glm-4-flash, glm-4-plus
    messages=[    
        {"role": "system", "content": system},    
        {"role": "user", "content": input} 
    ],
    top_p=0.7,
    temperature=0.9
 ) 

result_text = completion.choices[0].message.content
print(result_text)
len_text = len(result_text.split('\n'))
save_file.write(result_text)
